import os

import chromadb
from dotenv import load_dotenv
from langchain_chroma import Chroma
from langchain_core.messages import SystemMessage, HumanMessage

from lc_frame.lc_models.api import llm_api
from lc_frame.utils.log import logger

from langgraph.prebuilt import create_react_agent
from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreRouterToolkit

import warnings

warnings.filterwarnings('ignore')

current_path = os.path.abspath(__file__)
root_path = os.path.dirname(os.path.dirname(current_path))

pdfs_dict_path = os.path.join(root_path, "datas", "pdfs")
vector_db_path = os.path.join(root_path, "datas", "vector_db", "chroma")

# 加载key
dotenv_path = os.path.join(root_path, "lc_models", "api", ".qwen")
load_dotenv(dotenv_path=dotenv_path)


def get_current_datatime() -> str:
    """
    获取当前时间
    """
    from datetime import datetime
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def lc_rag(question=None):
    llm, _, embed = llm_api.get_qwen()
    vectorstore = None
    embeddings = llm_api.get_bge_embeddings()

    client = chromadb.PersistentClient(path=vector_db_path)
    if client.list_collections() != [] and client.get_collection(name="langchain") is not None:
        vectorstore = Chroma(embedding_function=embeddings,
                             client=client)
        logger.info("langchain 已存在，将使用 Chroma 进行文本搜索")

    toolkit = VectorStoreRouterToolkit(llm=llm, vectorstores=[{
        "vectorstore": vectorstore,
        "description": "这是一个chroma的向量数据库",
        "name": "langchain"
    }])

    tools = toolkit.get_tools()

    SQL_PREFIX = """
    请根据用户从私有知识库检索出来的上下文来回答用户的问题！
        请注意：
            1，如果用户的问题不在上下文中，请直接回答不知道！
            2，不要做任何解释，直接输出最终的结果即可！
    """
    system_message = SystemMessage(content=SQL_PREFIX)

    agent_executor = create_react_agent(model=llm, tools=tools, state_modifier=system_message)

    events = agent_executor.stream(
        {"messages": [HumanMessage(content=question)]},
        stream_mode="values"
    )

    result_list = []
    for event in events:
        result_list.append(event["messages"][-1].content)
        logger.info(event["messages"][-1].pretty_print())
    # 返回最终结果
    final_result = event["messages"][-1].content if result_list else None
    logger.info(f'最终答案是：{final_result}')
    return final_result


if __name__ == "__main__":
    question = "2007年度，工业业务和贸易业务分别占联化科技股份有限公司毛利比重为多少？"
    lc_rag(question)
